#!/usr/bin/env python3

import re
import os
import sys
import subprocess

# Sample kernel code input (replace this with your actual kernel code)
kernel_code = """
# Your entire kernel code here (the provided code)
"""
if len(sys.argv) == 2:
    with open(sys.argv[1], 'r') as f:
        kernel_code = f.read()

# Regex pattern to match kernel definitions
kernel_pattern = re.compile(
    r"(?P<kernel_name>\w+)\s*=\s*async_compile\.triton\('(?P<name>[\w_]+)',\s*'''(?P<code>.*?)'''",
    re.DOTALL
)

# Directory to save the extracted kernel files
output_dir = os.path.dirname(os.path.abspath(__file__))
os.makedirs(output_dir, exist_ok=True)

# Find all kernels in the kernel code
for match in kernel_pattern.finditer(kernel_code):
    kernel_name = match.group('name')
    kernel_body = match.group('code').strip()

    # Create the filename
    filename = os.path.join(output_dir, f"{kernel_name}.py")

    # Write the kernel code to the file
    with open(filename, 'w') as f:
        # f.write(f"{kernel_name} = async_compile.triton('{kernel_name}', '''{kernel_body}''', device_str='npu')\n")
        f.write(f"#!/usr/bin/env python3\n\n\n{kernel_body}\n")

    subprocess.run(["autopep8", "-i", filename])

print(f"Extracted kernels saved in '{output_dir}' directory.")
